{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "name": "Punctuation_and_Capitalization.ipynb",
   "provenance": [],
   "private_outputs": true,
   "collapsed_sections": [
    "daYw_Xll2ZR9"
   ]
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "accelerator": "GPU",
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "source": [],
    "metadata": {
     "collapsed": false
    }
   }
  }
 },
 "cells": [
    {
     "cell_type": "code",
     "execution_count": null,
     "metadata": {},
     "outputs": [],
     "source": [
     "BRANCH = 'r1.0.0rc1'"
     ]
    },
  {
   "cell_type": "code",
   "metadata": {
    "id": "o_0K1lsW1dj9",
    "colab_type": "code",
    "colab": {}
   },
   "source": [
    "\"\"\"\n",
    "You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
    "\n",
    "Instructions for setting up Colab are as follows:\n",
    "1. Open a new Python 3 notebook.\n",
    "2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",
    "3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",
    "4. Run this cell to set up dependencies.\n",
    "\"\"\"\n",
    "# If you're using Google Colab and not running locally, run this cell\n",
    "\n",
    "# install NeMo\n",
    "!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[nlp]\n"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "id": "uglDB-pVh__t",
    "colab_type": "code",
    "colab": {}
   },
   "source": [
    "# If you're not using Colab, you might need to upgrade jupyter notebook to avoid the following error:\n",
    "# 'ImportError: IProgress not found. Please update jupyter and ipywidgets.'\n",
    "\n",
    "! pip install ipywidgets\n",
    "! jupyter nbextension enable --py widgetsnbextension\n",
    "\n",
    "# Please restart the kernel after running this cell"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "dzqD2WDFOIN-",
    "colab_type": "code",
    "colab": {}
   },
   "source": [
    "from nemo.utils.exp_manager import exp_manager\n",
    "from nemo.collections import nlp as nemo_nlp\n",
    "\n",
    "import os\n",
    "import wget \n",
    "import torch\n",
    "import pytorch_lightning as pl\n",
    "from omegaconf import OmegaConf"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "daYw_Xll2ZR9",
    "colab_type": "text"
   },
   "source": [
    "# Task Description\n",
    "Automatic Speech Recognition (ASR) systems typically generate text with no punctuation and capitalization of the words. \n",
    "This tutorial explains how to implement a model in NeMo that will predict punctuation and capitalization for each word in a sentence to make ASR output more readable and to boost performance of the named entity recognition, machine translation or text-to-speech models. \n",
    "We'll show how to train a model for this task using a pre-trained BERT model. \n",
    "For every word in our training dataset we’re going to predict:\n",
    "\n",
    "- punctuation mark that should follow the word and\n",
    "- whether the word should be capitalized"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZnuziSwJ1yEB",
    "colab_type": "text"
   },
   "source": [
    "# Dataset\n",
    "This model can work with any dataset as long as it follows the format specified below. \n",
    "The training and evaluation data is divided into *2 files: text.txt and labels.txt*. \n",
    "Each line of the **text.txt** file contains text sequences, where words are separated with spaces: [WORD] [SPACE] [WORD] [SPACE] [WORD], for example:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TXFORGBv2Jqu",
    "colab_type": "text"
   },
   "source": [
    "\n",
    "\n",
    "```\n",
    "when is the next flight to new york\n",
    "the next flight is ...\n",
    "...\n",
    "```\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Y7ci55rM2QH8",
    "colab_type": "text"
   },
   "source": [
    "The **labels.txt** file contains corresponding labels for each word in text.txt, the labels are separated with spaces. Each label in labels.txt file consists of 2 symbols:\n",
    "\n",
    "- the first symbol of the label indicates what punctuation mark should follow the word (where O means no punctuation needed);\n",
    "- the second symbol determines if a word needs to be capitalized or not (where U indicates that the word should be upper cased, and O - no capitalization needed.)\n",
    "\n",
    "In this tutorial, we are considering only commas, periods, and question marks the rest punctuation marks were removed. To use more punctuation marks, update the dataset to include desired labels, no changes to the model needed. \n",
    "\n",
    "Each line of the **labels.txt** should follow the format: \n",
    "[LABEL] [SPACE] [LABEL] [SPACE] [LABEL] (for labels.txt). \n",
    "For example, labels for the above text.txt file should be:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-94C1-864EW1",
    "colab_type": "text"
   },
   "source": [
    "\n",
    "\n",
    "```\n",
    "OU OO OO OO OO OO OU ?U\n",
    "OU OO OO OO ...\n",
    "...\n",
    "```\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VsEmwIPO4L4V",
    "colab_type": "text"
   },
   "source": [
    "The complete list of all possible labels for this task used in this tutorial is: `OO, ,O, .O, ?O, OU, ,U, .U, ?U.`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SL58EWkd2ZVb",
    "colab_type": "text"
   },
   "source": [
    "## Download and preprocess the data¶"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "THi6s1Qx2G1k",
    "colab_type": "text"
   },
   "source": [
    "In this notebook we are going to use a subset of English examples from the [Tatoeba collection of sentences](https://tatoeba.org/eng) this script will download and preprocess the Tatoeba data [NeMo/examples/nlp/token_classification/get_tatoeba_data.py](https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/token_classification/data/get_tatoeba_data.py). Note, for further experiments with the model, set NUM_SAMPLES=-1 and consider including other datasets to improve model performance. \n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "n8HZrDmr12_-",
    "colab_type": "code",
    "colab": {}
   },
   "source": [
    "DATA_DIR = \"PATH_TO_DATA\"\n",
    "WORK_DIR = \"PATH_TO_CHECKPOINTS_AND_LOGS\"\n",
    "MODEL_CONFIG = \"punctuation_capitalization_config.yaml\"\n",
    "\n",
    "# model parameters\n",
    "BATCH_SIZE = 128\n",
    "MAX_SEQ_LENGTH = 64\n",
    "LEARNING_RATE = 0.00002\n",
    "NUM_SAMPLES = 10000"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "UOeeeCGqI-9c",
    "colab_type": "code",
    "colab": {}
   },
   "source": [
    "## download get_tatoeba_data.py script to download and preprocess the Tatoeba data\n",
    "os.makedirs(WORK_DIR, exist_ok=True)\n",
    "if not os.path.exists(WORK_DIR + '/get_tatoeba_data.py'):\n",
    "    print('Downloading get_tatoeba_data.py...')\n",
    "    wget.download(f'https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/examples/nlp/token_classification/data/get_tatoeba_data.py', WORK_DIR)\n",
    "else:\n",
    "    print ('get_tatoeba_data.py is already exists')"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "H0ulD2TL13DR",
    "colab_type": "code",
    "colab": {}
   },
   "source": [
    "# download and preprocess the data\n",
    "# --clean_dir flag deletes raw Tataoeba data, remove the flag to avoid multiple data downloads if you want to experiment with the dataset size\n",
    "! python $WORK_DIR/get_tatoeba_data.py --data_dir $DATA_DIR --num_sample $NUM_SAMPLES --clean_dir"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pctMm2vsPlPT",
    "colab_type": "text"
   },
   "source": [
    "after execution of the above cell, your data folder will contain the following 4 files needed for training (raw Tatoeba data could be present if `--clean_dir` was not used):\n",
    "- labels_dev.txt\n",
    "- labels_train.txt\n",
    "- text_dev.txt\n",
    "- text_train.txt\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "KKwxXXezPvXF",
    "colab_type": "code",
    "colab": {}
   },
   "source": [
    "! ls -l {DATA_DIR}"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "6UDPgadLN6SG",
    "colab_type": "code",
    "colab": {}
   },
   "source": [
    "# let's take a look at the data \n",
    "print('Text:')\n",
    "! head -n 5 {DATA_DIR}/text_train.txt\n",
    "\n",
    "print('\\nLabels:')\n",
    "! head -n 5 {DATA_DIR}/labels_train.txt"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "daludzzL2Jba",
    "colab_type": "text"
   },
   "source": [
    "# Model Configuration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_whKCxfTMo6Y",
    "colab_type": "text"
   },
   "source": [
    "In the Punctuation and Capitalization Model, we are jointly training two token-level classifiers on top of the pretrained [BERT](https://arxiv.org/pdf/1810.04805.pdf) model: \n",
    "- one classifier to predict punctuation and\n",
    "- the other one - capitalization.\n",
    "\n",
    "The model is defined in a config file which declares multiple important sections. They are:\n",
    "- **model**: All arguments that are related to the Model - language model, token classifiers, optimizer and schedulers, dataset and any other related information\n",
    "\n",
    "- **trainer**: Any argument to be passed to PyTorch Lightning"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "T1gA8PsJ13MJ",
    "colab_type": "code",
    "colab": {}
   },
   "source": [
    "# download the model's configuration file \n",
    "config_dir = WORK_DIR + '/configs/'\n",
    "os.makedirs(config_dir, exist_ok=True)\n",
    "if not os.path.exists(config_dir + MODEL_CONFIG):\n",
    "    print('Downloading config file...')\n",
    "    wget.download(f'https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/examples/nlp/token_classification/conf/' + MODEL_CONFIG, config_dir)\n",
    "else:\n",
    "    print ('config file is already exists')"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "mX3KmWMvSUQw",
    "colab_type": "code",
    "colab": {}
   },
   "source": [
    "# this line will print the entire config of the model\n",
    "config_path = f'{WORK_DIR}/configs/{MODEL_CONFIG}'\n",
    "print(config_path)\n",
    "config = OmegaConf.load(config_path)\n",
    "print(OmegaConf.to_yaml(config))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZCgWzNBkaQLZ",
    "colab_type": "text"
   },
   "source": [
    "# Setting up Data within the config\n",
    "\n",
    "Among other things, the config file contains dictionaries called dataset, train_ds and validation_ds. These are configurations used to setup the Dataset and DataLoaders of the corresponding config.\n",
    "\n",
    "If both training and evaluation files are located in the same directory, simply specify `model.dataset.data_dir`, like we are going to do below.\n",
    "However, if your evaluation files are located in a different directory, or you want to use multiple datasets for evaluation, specify paths to the directory(ies) with evaluation file(s) in the following way:\n",
    "\n",
    "`model.validation_ds.ds_item=[PATH_TO_DEV1,PATH_TO_DEV2]` (Note no space between the paths and square brackets).\n",
    "\n",
    "Also notice that some configs, including `model.dataset.data_dir`, have `???` in place of paths, this values are required to be specified by the user.\n",
    "\n",
    "Let's now add the data directory path to the config."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "LQHCJN-ZaoLp",
    "colab_type": "code",
    "colab": {}
   },
   "source": [
    "# in this tutorial train and dev data is located in the same folder, so it is enought to add the path of the data directory to our config\n",
    "config.model.dataset.data_dir = DATA_DIR"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nB96-3sTc3yk",
    "colab_type": "text"
   },
   "source": [
    "# Building the PyTorch Lightning Trainer\n",
    "\n",
    "NeMo models are primarily PyTorch Lightning modules - and therefore are entirely compatible with the PyTorch Lightning ecosystem!\n",
    "\n",
    "Let's first instantiate a Trainer object!"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "1tG4FzZ4Ui60",
    "colab_type": "code",
    "colab": {}
   },
   "source": [
    "print(\"Trainer config - \\n\")\n",
    "print(OmegaConf.to_yaml(config.trainer))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "knF6QeQQdMrH",
    "colab_type": "code",
    "colab": {}
   },
   "source": [
    "# lets modify some trainer configs\n",
    "# checks if we have GPU available and uses it\n",
    "cuda = 1 if torch.cuda.is_available() else 0\n",
    "config.trainer.gpus = cuda\n",
    "config.trainer.precision = 16 if torch.cuda.is_available() else 32\n",
    "\n",
    "# For mixed precision training, use precision=16 and amp_level=O1\n",
    "\n",
    "# Reduces maximum number of epochs to 1 for a quick training\n",
    "config.trainer.max_epochs = 1\n",
    "\n",
    "# Remove distributed training flags\n",
    "config.trainer.accelerator = None\n",
    "\n",
    "trainer = pl.Trainer(**config.trainer)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8IlEMdVxdr6p",
    "colab_type": "text"
   },
   "source": [
    "# Setting up a NeMo Experiment¶\n",
    "\n",
    "NeMo has an experiment manager that handles logging and checkpointing for us, so let's use it!"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "8uztqGAmdrYt",
    "colab_type": "code",
    "colab": {}
   },
   "source": [
    "exp_dir = exp_manager(trainer, config.get(\"exp_manager\", None))\n",
    "\n",
    "# the exp_dir provides a path to the current experiment for easy access\n",
    "exp_dir = str(exp_dir)\n",
    "exp_dir"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6FI_nQsJo_11",
    "colab_type": "text"
   },
   "source": [
    "# Model Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8tjLhUvL_o7_",
    "colab_type": "text"
   },
   "source": [
    "Before initializing the model, we might want to modify some of the model configs. For example, we might want to modify the pretrained BERT model."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "Xeuc2i7Y_nP5",
    "colab_type": "code",
    "colab": {}
   },
   "source": [
    "# complete list of supported BERT-like models\n",
    "nemo_nlp.modules.get_pretrained_lm_models_list()\n",
    "\n",
    "PRETRAINED_BERT_MODEL = \"bert-base-uncased\""
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "RK2xglXyAUOO",
    "colab_type": "code",
    "colab": {}
   },
   "source": [
    "# add the specified above model parameters to the config\n",
    "config.model.language_model.pretrained_model_name = PRETRAINED_BERT_MODEL\n",
    "config.model.train_ds.batch_size = BATCH_SIZE\n",
    "config.model.validation_ds.batch_size = BATCH_SIZE\n",
    "config.model.optim.lr = LEARNING_RATE\n",
    "config.model.train_ds.num_samples = NUM_SAMPLES\n",
    "config.model.validation_ds.num_samples = NUM_SAMPLES\n"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "id": "gYKcziSsiAAd",
    "colab_type": "text"
   },
   "source": [
    "Now, we are ready to initialize our model. During the model initialization call, the dataset and data loaders we'll be prepared for training and evaluation.\n",
    "Also, the pretrained BERT model will be downloaded, note it can take up to a few minutes depending on the size of the chosen BERT model."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "id": "Yk2hJssviAAe",
    "colab_type": "code",
    "colab": {}
   },
   "source": [
    "# initialize the model\n",
    "# during this stage, the dataset and data loaders we'll be prepared for training and evaluation\n",
    "config.trainer.max_epochs = 3\n",
    "model = nemo_nlp.models.PunctuationCapitalizationModel(cfg=config.model, trainer=trainer)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kQ592Tx4pzyB",
    "colab_type": "text"
   },
   "source": [
    "## Monitoring training progress\n",
    "Optionally, you can create a Tensorboard visualization to monitor training progress."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "mTJr16_pp0aS",
    "colab_type": "code",
    "colab": {}
   },
   "source": [
    "try:\n",
    "  from google import colab\n",
    "  COLAB_ENV = True\n",
    "except (ImportError, ModuleNotFoundError):\n",
    "  COLAB_ENV = False\n",
    "\n",
    "# Load the TensorBoard notebook extension\n",
    "if COLAB_ENV:\n",
    "  %load_ext tensorboard\n",
    "  %tensorboard --logdir {exp_dir}\n",
    "else:\n",
    "  print(\"To use tensorboard, please use this notebook in a Google Colab environment.\")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "hUvnSpyjp0Dh",
    "colab_type": "code",
    "colab": {}
   },
   "source": [
    "# start the training\n",
    "trainer.fit(model)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VPdzJVAgSFaJ",
    "colab_type": "text"
   },
   "source": [
    "# Inference\n",
    "\n",
    "To see how the model performs, let’s run inference on a few examples."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "DQhsamclRtxJ",
    "colab_type": "code",
    "colab": {}
   },
   "source": [
    "# define the list of queiries for inference\n",
    "queries = [\n",
    "        'we bought four shirts and one mug from the nvidia gear store in santa clara',\n",
    "        'what can i do for you today',\n",
    "        'how are you',\n",
    "        'how is the weather in',\n",
    "    ]\n",
    "inference_results = model.add_punctuation_capitalization(queries)\n",
    "print()\n",
    "\n",
    "for query, result in zip(queries, inference_results):\n",
    "    print(f'Query   : {query}')\n",
    "    print(f'Combined: {result.strip()}\\n')"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ref1qSonGNhP",
    "colab_type": "text"
   },
   "source": [
    "## Training Script\n",
    "\n",
    "If you have NeMo installed locally, you can also train the model with [nlp/token_classification/punctuation_capitalization_train.py](https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/token_classification/punctuation_capitalization_train.py).\n",
    "\n",
    "To run training script, use:\n",
    "\n",
    "`python punctuation_and_capitalization_train.py model.dataset.data_dir=PATH_TO_DATA_DIR`\n",
    "\n",
    "Set NUM_SAMPLES=-1 and consider including other datasets to improve the performance of the model.\n",
    "\n",
    "# Finetuning model with your data\n",
    "\n",
    "When we were training the model from scratch, the datasets were prepared for training during the model initialization. When we are using a pretrained Punctuation and Capitalization model, before training, we need to setup training and evaluation data."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "4X1BahRlkaNf",
    "colab_type": "code",
    "colab": {}
   },
   "source": [
    "# let's reload our pretrained model\n",
    "pretrained_model = nemo_nlp.models.PunctuationCapitalizationModel.from_pretrained('punctuation_en_distilbert')\n",
    "\n",
    "# setup train and validation Pytorch DataLoaders\n",
    "pretrained_model.update_data_dir(DATA_DIR)\n",
    "pretrained_model.setup_training_data()\n",
    "pretrained_model.setup_validation_data()\n",
    "\n",
    "# and now we can create a PyTorch Lightning trainer and call `fit` again\n",
    "# for this tutorial we are setting fast_dev_run to True, and the trainer will run 1 training batch and 1 validation batch\n",
    "# for actual model training, disable the flag\n",
    "fast_dev_run = True\n",
    "trainer = pl.Trainer(gpus=1, fast_dev_run=fast_dev_run)\n",
    "trainer.fit(pretrained_model)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "l7A5FeiTl6Zd",
    "colab_type": "code",
    "colab": {}
   },
   "source": [
    ""
   ],
   "execution_count": null,
   "outputs": []
  }
 ]
}